#!/usr/bin/env python
import random

import numpy as np
import matplotlib.pyplot as plt

from problem import Poisson_1d, Helmholtz_1d, Allen_cahn_1d, Boundary_layer_1d, Poisson_1d_complex
from options import Options
from regression import LinearRegression
from feature_extractor import Gauss1DFeatureExtractor, Gauss2DFeatureExtractor
from utils import add_axes
from itertools import accumulate


class Solver:
    def __init__(self, args, print_info=True):
        self.domain = args.domain
        self.a, self.b = self.domain
        self.problem_name = args.problem_name
        self.case = args.case
        self.k = args.k
        self.n_gaussian = args.n_gaussian
        self.n_level = args.n_level
        self.n_pde = args.n_pde
        self.sigma = args.sigma
        self.mode = args.mode

        # Set problem, extractor, regressor
        self.problem = self._problem()
        self.extractor = Gauss1DFeatureExtractor(self.n_gaussian, n_level=self.n_level, sigma=self.sigma, mode=self.mode, a=self.a, b=self.b)
        if print_info:
            print("Loading extractor")
            print(f"{self.extractor}")
        self.regressor = LinearRegression()

        # seed
        np.random.seed(1234)

        # Dataset
        self.x_pde = np.random.uniform(self.a, self.b, self.n_pde)    # (n_pde,) suitable for regressor
        self.x_bc = np.array([self.a, self.b])

    def _problem(self):
        assert self.problem_name in ["poisson", "helmholtz", "allen_cahn", "boundary_layer", "poisson_complex"], f"need select from [poisson, helmholtz, allen_cahn]"
        if self.problem_name == "poisson":
            return Poisson_1d(case=self.case, k=self.k, domain=self.domain)
        elif self.problem_name == "helmholtz":
            return Helmholtz_1d(case=self.case, k=self.k, domain=self.domain)
        elif self.problem_name == "allen_cahn":
            return Allen_cahn_1d(case=self.case, k=self.k, domain=self.domain)
        elif self.problem_name == "boundary_layer":
            return Boundary_layer_1d(case=self.case, k=self.k, domain=self.domain)
        elif self.problem_name == "poisson_complex":
            return Poisson_1d_complex(case=self.case, k=self.k, domain=self.domain)

    def get_features_target(self):
        if self.problem_name == "poisson":
            return self.get_features_target_poisson()
        elif self.problem_name == "helmholtz":
            return self.get_features_target_helmholtz()
        elif self.problem_name == "boundary_layer":
            return self.get_features_target_boundary_layer()
        elif self.problem_name == "poisson_complex":
            return self.get_features_target_poisson_complex()

    def get_features_target_poisson(self):
        '''
        u_xx = f
        '''
        z_pde = self.extractor.eval(self.x_pde, eval_list=('u00',))
        z_bc = self.extractor.eval(self.x_bc, eval_list=('u',))

        features_pde = z_pde['u00']
        features_bc = z_bc['u']

        features = np.vstack([features_pde, features_bc])

        target_pde = self.problem.rhs(self.x_pde)
        target_bc = self.problem.bc(self.x_bc)
        target = np.hstack([target_pde, target_bc])

        return features, target
    
    def get_features_target_poisson_complex(self):
        '''
        u_xx = f
        '''
        z_pde = self.extractor.eval(self.x_pde, eval_list=('u0','u00'))
        z_bc = self.extractor.eval(self.x_bc, eval_list=('u',))

        a = self.problem.a(self.x_pde)
        a_x = self.problem.a_x(self.x_pde)
        features_pde = - a[:, np.newaxis] * z_pde['u00'] - a_x[:, np.newaxis] * z_pde['u0']
        features_bc = z_bc['u']

        features = np.vstack([features_pde, features_bc])

        target_pde = self.problem.rhs(self.x_pde)
        target_bc = self.problem.bc(self.x_bc)
        target = np.hstack([target_pde, target_bc])

        return features, target

    def get_features_target_helmholtz(self):
        z_pde = self.extractor.eval(self.x_pde, eval_list=('u00', 'u'))
        z_bc = self.extractor.eval(self.x_bc, eval_list=('u',))

        features_pde = z_pde['u00'] + self.k * self.k * z_pde['u']
        features_bc = z_bc['u']

        features = np.vstack([features_pde, features_bc])

        target_pde = self.problem.rhs(self.x_pde)
        target_bc = self.problem.bc(self.x_bc)
        target = np.hstack([target_pde, target_bc])

        return features, target

    def get_features_target_boundary_layer(self):
        z_pde = self.extractor.eval(self.x_pde, eval_list=('u00', 'u0'))
        z_bc = self.extractor.eval(self.x_bc, eval_list=('u',))

        features_pde = z_pde['u0'] - self.k * z_pde['u00']
        features_bc = z_bc['u']

        features = np.vstack([features_pde, features_bc])

        target_pde = self.problem.rhs(self.x_pde)
        target_bc = self.problem.bc(self.x_bc)
        target = np.hstack([target_pde, target_bc])

        return features, target
    
    def get_NTK_matrix(self):
        '''
           du/dw 构造Jacobi矩阵J(N_pde x N_w)
           下面只考虑到possion方程
        '''
        def normalize_matrix(matrix):
            A_min = np.min(matrix)
            A_max = np.max(matrix)
            A_norm = 2 * (matrix - A_min) / (A_max - A_min) - 1
            return A_norm

        test_x_pde = np.linspace(self.a, self.b, 100)
        z_pde = self.extractor.eval(test_x_pde, eval_list=('u00', ))
        J = z_pde['u00']
        ntk_matrix = J @ J.T
        normal_ntk = normalize_matrix(ntk_matrix)   # 归一化到[-1, 1]
        plt.imshow(normal_ntk, cmap='coolwarm', vmin=-1, vmax=1)   # 'viridis', 'rainbow', 'coolwarm'
        plt.colorbar()
        plt.show()
        return normal_ntk
    
    def plot_gauss_shape(self, ax, level=5):
        '''
           绘制最后一层高斯函数的形状
        '''
        if level > self.extractor.n_level or self.extractor.info[level]["points_num"] == 0:
            return
        center_arr = np.empty(0)
        sigma_arr = np.empty(0)
        num_gs = [group["points_num"] for group in self.extractor.info if group]
        accum_num = [0] + list(accumulate(num_gs))
        for l in range(level, level + 1):   # (1, self.extractor.n_level + 1)
            center = self.extractor.info[l]["coordinates"]
            sigma = self.extractor.info[l]["sigma"]
            n_gauss = self.extractor.info[l]["points_num"]
            center_arr = np.concatenate([center_arr, center])
            sigma_arr = np.concatenate([sigma_arr, np.array([sigma]*n_gauss)])

        x = np.linspace(self.a, self.b, 500)

        x_expanded = x[:, np.newaxis]
        means_expanded = center_arr[np.newaxis, :]
        stds_expanded = sigma_arr[np.newaxis, :]
        
        exponent = -((x_expanded - means_expanded) ** 2) / (2 * stds_expanded ** 2)

        # 不带权重
        # pdf = np.exp(exponent) * (2*np.random.random((1, exponent.shape[1])) - 1)   # 给一个随机的权重
        # ax.plot(x_expanded, pdf)
        # ax.set_title("Gauss shape (without weight)")

        # 所有Gauss的形状
        # pdf = np.exp(exponent) * self.regressor.w[np.newaxis, accum_num[level-1]:accum_num[level]]
        # ax.plot(x_expanded, pdf)
        # ax.set_title(f"Gauss shape (level:{level})")

        # 将同一层级的Gauss函数进行求和
        pdf = np.exp(exponent) * self.regressor.w[np.newaxis, accum_num[level-1]:accum_num[level]]
        sum_gauss = np.sum(pdf, axis=1)
        ax.plot(x_expanded, sum_gauss)
        ax.set_title(f"Gauss sum (level:{level})")

    def plot_w(self, ax, level=None):
        '''
           绘制出某层级回归器去求解出来的权重w
        '''
        if level:
            num_gs = [group["points_num"] for group in self.extractor.info if group]
            accum_num = [0] + list(accumulate(num_gs))
            begin_index, end_index = accum_num[level-1], accum_num[level] 
            ax.scatter(range(begin_index,end_index),self.regressor.w[begin_index:end_index])
        else:
            # 绘制所有
            n = len(self.regressor.w)
            ax.scatter(range(n), self.regressor.w)
            ax.set_title("w")

    def delete(self):
        '''
           删点多余的层级
        '''
        t1, t2 = 0.01, 0.1
        # 依次检验2-max_levle 层
        num_gs = [group["points_num"] for group in self.extractor.info if group]
        accum_num = [0] + list(accumulate(num_gs))
        x = np.linspace(self.a, self.b, 500)   # 采样点的数值进行检测
        x_expanded = x[:, np.newaxis]
        record = None     # 保留上一层的记录
        for l in range(1, self.extractor.n_level + 1):   # (1, self.extractor.n_level + 1)
            center_arr = np.empty(0)
            sigma_arr = np.empty(0)
            center = self.extractor.info[l]["coordinates"]
            sigma = self.extractor.info[l]["sigma"]
            n_gauss = self.extractor.info[l]["points_num"]
            if n_gauss == 0:
                continue
            center_arr = np.concatenate([center_arr, center])
            sigma_arr = np.concatenate([sigma_arr, np.array([sigma]*n_gauss)])
            means_expanded = center_arr[np.newaxis, :]
            stds_expanded = sigma_arr[np.newaxis, :]
            exponent = -((x_expanded - means_expanded) ** 2) / (2 * stds_expanded ** 2)
            pdf = np.exp(exponent) * self.regressor.w[np.newaxis, accum_num[l-1]:accum_num[l]]
            sum_gauss = np.sum(pdf, axis=1)
            # 通过层级权重过小进行删除
            big_val = np.max(np.abs(sum_gauss))
            if big_val < t1:
                self.extractor.deleteLevel(l)
                print(f"Level {l} has too small weight {big_val}, delete")
            sum_gauss = (sum_gauss - np.min(sum_gauss)) / (np.max(sum_gauss) - np.min(sum_gauss) + 1e-5)  # 归一化
            # 通过两层之间的误差进行删除
            if record is not None:
                error = np.max(np.abs(sum_gauss - record))
                if error < t2:
                    self.extractor.deleteLevel(l)
                    print(f"error is {error:.2e}, Level {l} has the same shape as before level, delete")
                else:
                    print(f"error is {error:.2e}, Level {l} does not delete")
            record = sum_gauss

    
    def update(self):
        num_gs = [group["points_num"] for group in self.extractor.info if group]
        accum_num = [0] + list(accumulate(num_gs))
        
        for level in range(3, self.extractor.n_level + 1):  
            if self.extractor.info[level]["points_num"] == 0:
                continue
            centers = self.extractor.info[level]["coordinates"]
            sigma = self.extractor.info[level]["sigma"]
            level_w = np.abs(self.regressor.w[accum_num[level-1]:accum_num[level]])

            # dist_matrix = centers[:, np.newaxis] - centers

            # # 计算综合相似性矩阵
            # similarity_matrix = np.exp(-dist_matrix ** 2) 

            # # 计算相似性加权贡献比矩阵
            # weighted_contribution_ratio = similarity_matrix / level_w[:, np.newaxis]
            # redundancy_scores = np.mean(weighted_contribution_ratio, axis=1)
            # indices = redundancy_scores > np.percentile(redundancy_scores, 10)

            indices = level_w > np.maximum(0.001, np.percentile(level_w, 10))
            self.extractor.updateLevel(level, indices)


        

    def solve(self, plot=True, print_info=True):
        # nonlinear iteration
        if self.problem_name == "allen_cahn":
            return self.picard_allen_cahn(plot=plot)

        # 1. features and target
        features, target = self.get_features_target()

        # 2. Least square
        self.regressor.fit(features, target)

        # 3. Test
        x = np.linspace(self.a, self.b, 10000)
        u = self.problem.u_exact(x)

        z = self.extractor.eval(x, eval_list=('u'))
        features = z['u']

        u_pred = self.regressor.predict(features)
        mse = np.sqrt(np.mean((u - u_pred) ** 2))
        mae = np.mean(np.abs(u- u_pred))

        if print_info:
            print(f"problem_name: {self.problem_name}, case: {self.case}, Gaussians: {self.n_gaussian:4d}, "
              f"sigma: {self.sigma:.3f}, RMSE: {mse:.2e}, MAE: {mae:.2e}, "
              f"SumGaussians:{self.extractor.getSumGaussians()}")

        if plot:
            fig = plt.figure(figsize=(8, 8))

            ax = fig.add_subplot(111)
            # ax.plot(x, u, 'k-', x, u_pred, 'r-')
            ax.plot(x, u, 'k-', label="ground truth")
            ax.plot(x, u_pred, 'r--', label="predict")
            # add_axes(ax, x, u_pred, u, 0.8, 0.81, -1.3, 1.3) #添加放大区域
            ax.set_title("solution comparison")
            ax.legend()

            # ax = fig.add_subplot(122)
            # ax.plot(x, np.abs(u_pred - u))
            # ax.set_title("MAE error")

            # ax = fig.add_subplot(122)
            # self.plot_gauss_shape(ax, level=1)

            # ax = fig.add_subplot(233)
            # self.plot_gauss_shape(ax, level=2)

            # ax = fig.add_subplot(234)
            # self.plot_gauss_shape(ax, level=3)

            # ax = fig.add_subplot(235)
            # self.plot_gauss_shape(ax, level=4)

            # ax = fig.add_subplot(236)
            # self.plot_gauss_shape(ax, level=5)
            plt.show()

        return mse, mae

    def picard_allen_cahn(self, plot=True):
        """
            nonlinear picard
            just for allen-cahn equation
        """
        # initial guess
        u_pde_old = np.zeros((self.n_pde, 1))    # dim=1

        # 1. features and target
        z_pde = self.extractor.eval(self.x_pde, eval_list=('u00', 'u'))
        z_bc = self.extractor.eval(self.x_bc, eval_list=('u',))

        target_pde = self.problem.rhs(self.x_pde)
        target_bc = self.problem.bc(self.x_bc)
        target = np.hstack([target_pde, target_bc])

        # 2. Least square self.regressor

        # 3. Picard
        print("Start picard iteration")
        print("Iteration   res")
        eps, iteration, res, max_iter = 1e-5, 0, np.inf, 20
        while iteration < max_iter and res > eps:
            iteration += 1

            prev_cof = u_pde_old * u_pde_old - 1

            features_pde = z_pde['u00'] + prev_cof * z_pde['u']
            features_bc = z_bc['u']

            features = np.vstack([features_pde, features_bc])

            self.regressor.fit(features, target)

            u_pde_new = self.regressor.predict(z_pde['u'])
            res = np.linalg.norm(u_pde_new - u_pde_old.flatten())
            u_pde_old = u_pde_new[:, np.newaxis]
            print(f"{iteration:6d}  {res:10.3e}")

        # 4.Test
        x = np.linspace(self.a, self.b, 2000)
        u = self.problem.u_exact(x)

        z = self.extractor.eval(x, eval_list=('u'))
        features = z['u']

        u_pred = self.regressor.predict(features)
        rmse = np.sqrt(np.mean((u - u_pred) ** 2))
        mae = np.mean(np.abs(u- u_pred))

        print(f"problem_name: {self.problem_name}, case: {self.case}, Gaussians: {self.n_gaussian:4d}, "
              f"sigma: {self.sigma:.3f}, RMSE: {rmse:.2e}, MAE:{mae:.2e}"
              f"SumGaussians:{self.extractor.getSumGaussians()}")

        if plot:
            fig = plt.figure(figsize=(12, 5))

            # ax = fig.add_subplot(121)
            # ax.plot(self.regressor.w, label="parameters")
            # ax.legend()

            ax = fig.add_subplot(121)
            ax.plot(x, np.abs(u - u_pred))

            ax = fig.add_subplot(122)
            ax.plot(x, u, 'k-', label="ground truth")
            ax.plot(x, u_pred, 'r-', label="predict")
            ax.legend()

            plt.show()

        return rmse, mae


if __name__ == '__main__':
    args = Options().parse()
    args.domain = (0, 1) # (100, 101)
    args.problem_name = 'allen_cahn'  # ['poisson, helmholtz, allen_cahn', 'boundary_layer', 'poisson_complex']
    args.case = 2
    args.k = 100      # some problem case will use  np.sqrt(9200)  0.005  0.05
    args.n_gaussian = 25    # sigma_1
    args.n_level = 5
    args.n_pde = 2000
    args.sigma = 0.1
    # args.mode = 1
    print("#"*10 + "args information" + "#"*10)
    print(args)

    solver = Solver(args)
    solver.solve()

    # print(f"k   n_level    rmse   mae")
    # for args.k in [100, 300, 500, 1000, 1500]:
    #     for args.n_level in [5, 6, 7]:
    #         solver = Solver(args, print_info=False)
    #         rmse, mae = solver.solve(plot=False, print_info=False)
    #         print(f"{args.k}, {args.n_level}, {rmse:.2e}, {mae:.2e}")

    # 优化压缩 改变solver

    # for _ in range(5):
    #     solver.delete()
    #     solver.update()
    #     print("update extractor")
    #     print(solver.extractor)
    #     solver.solve()

    # fig = plt.figure(figsize=(8, 8))
    # ax = fig.add_subplot(111)
    # x = np.linspace(0, 1, 2000)
    
    # for args.k in [1, 0.5, 0.1, 0.05, 0.01, 0.005]:
    #     solver = Solver(args)
    #     _, _, u = solver.solve(plot=False)
    #     ax.plot(x, u, label=f"{args.k}")
    # ax.legend()
    # ax.set_xlabel('x')
    # ax.set_ylabel('y')
    # plt.show()
    


